# -*- coding: utf-8 -*-
import base64
import json
from math import sqrt

import requests
import urllib

from tool import *

# 打开客户端端口配置文件

# 图像数据转换
# img_path = './picture_test/'
# save_path = './ResultJson/'
# img = cv2.imread(img_path)
# filelist = os.listdir(img_path)
# for filename in filelist:
#     print("************" + filename + "*****************")
#     with open(img_path + filename, 'rb') as f:
#         base64_data = base64.b64encode(f.read())
#     # 网络传输
#     res = {'base64img': base64_data, 'isAnalysis': 1}
#     r = requests.post(url, data=res)
#     print("url request Completed!")
#     result = str(r.text)
#     json_str = json.loads(result)
#     f2 = open(save_path + filename[:-4] + '.json', 'w', encoding="utf-8")
#     f2.write(json.dumps(json_str, ensure_ascii=False, indent=4, separators=(',', ':')))
#     f2.close()
# print("Complete image recognition! Result has been written!")


OCRSTATE = {2: '授权错误', 3: '图像数据错误', 4: '功能未实现', 5: '值为空'}


def getOcrState(code):
    if code in OCRSTATE.keys():
        return '[%d]%s' % (code, OCRSTATE[code])
    return '[%d]未知' % (code)


# 识别文件汉王
def ocrHanWangName(url, name):
    # BCD串生成
    base64_data = ''
    with open(name, 'rb') as f:
        base64_data = base64.b64encode(f.read()).decode("utf-8")
    # 网络传输
    res = {'base64img': base64_data, 'isAnalysis': 1}
    return requests.post(url, data=res).json()


# 识别文件百度
def ocrBaiduName(url, name, maxByte):
    # BCD串生成
    # 判断文件大小
    bytes = GetFileByte(name)
    if maxByte > 0 and bytes > maxByte:
        # 大于最大限制,生成新的可识别jpg图片
        pass

    base64_data = ''
    with open(name, "rb") as f:
        base64_data = base64.b64encode(f.read()).decode("utf8")
        base64_data = urllib.parse.quote_plus(base64_data)
    # 网络传输
    payload = 'image=%s&language_type=%s&detect_direction=%s&vertexes_location=%s&paragraph=%s&probability=%s' \
              % (base64_data, 'CHN_ENG', 'false', 'false', 'false', 'false')
    headers = {
        'Content-Type': 'application/x-www-form-urlencoded',
        'Accept': 'application/json'
    }
    response = requests.request("POST", url, headers=headers, data=payload)
    return response.json()


OCRTABLE = {'type': 2}
OCRIMAGE = {'type': 1, 'coords': [0, 0, 0, 0, 0, 0, 0, 0]}
OCRLINES = {'type': 0, 'code': '', 'coords': [0, 0, 0, 0, 0, 0, 0, 0], 'nums': 0,
            'script': 0, 'style': 0, 'area': 0, 'direction': 0, 'chars': []}
OCRCHARS = {'code': '', 'coord': [0, 0, 0, 0], 'score': 0, 'area': 0}


# 计算面积和判断方向
def getAreaDirection(nums, coords):
    area = 0
    direction = 0
    x1 = min(coords[0], coords[6])
    y1 = min(coords[1], coords[3])
    x2 = max(coords[2], coords[4])
    y2 = max(coords[5], coords[7])
    width = x2 - x1
    height = y2 - y1
    area = int(width * height / nums)
    b = int(sqrt(area))
    if nums > 1:
        if (height / b) > 2.0:
            direction = 1
    return area, direction


# 处理识别汉王结果
def ocrProcHanWang(name, ocrJson):
    # 每个单词列表 ['code':,]
    info = getImgFile(name)
    ocrDict = {'name': name, 'info': info, 'result': [], 'data': []}
    if ocrJson == None:
        # 图片，未识别
        return ocrDict
    ocrDict['result'] = ocrJson.copy()
    for l1 in ocrJson:
        # 判断每一组字典，可能多行
        # data-['type'-0文本/1表格，'code'-信息，'coord'-[左上,右上,右下,左下]，'nums'-字数，
        # 'script'-类型0手写/1印刷，'style'-类型0标题/1正文，'area'-单面积，'direction'-0横向/1纵向，'chars'-]
        # chars-['code'-字符，'coord'-[左,右,上,下]，'score'-分数0-1000，'area'-面积]
        if l1['comp_type'] == 0:
            # 文本行
            lineNums = l1['line_num']
            # 行数
            if lineNums <= 0:
                continue
            lineList = l1['line']
            if lineNums != len(lineList):
                continue
            for l2 in lineList:
                line = OCRLINES.copy()
                char = OCRCHARS.copy()
                line['type'] = 0
                line['code'] = l2['code']
                line['coords'] = l2['coords']
                line['nums'] = l2['char_num']
                line['script'] = l2['Script']
                line['style'] = l2['style']
                line['area'], line['direction'] = getAreaDirection(line['nums'], line['coords'])
                line['chars'] = []
                charList = l2['chars']
                for l3 in charList:
                    char['code'] = l3['code']
                    char['coord'] = l3['coords']
                    char['score'] = l3['score']
                    char['area'] = (char['coord'][1] - char['coord'][0]) * (char['coord'][3] - char['coord'][2])
                    line['chars'].append(char)
                ocrDict['data'].append(line)
        if l1['comp_type'] == 1:
            # 图像
            line = OCRIMAGE.copy()
            line['type'] = 1
            line['coords'] = l1['coords']
            ocrDict['data'].append(line)
            continue
        if l1['comp_type'] == 2:
            # 表格
            line = OCRTABLE.copy()
            line['type'] = 2
            ocrDict['data'].append(line)
            continue

    return ocrDict


# 处理识别百度结果
def ocrProcBaidu(name, ocrJson):
    # 每个单词列表 ['code':,]
    info = getImgFile(name)
    ocrDict = {'name': name, 'info': info, 'result': [], 'data': []}
    if ocrJson == None:
        # 图片，未识别
        return ocrDict
    ocrDict['result'] = ocrJson.copy()
    # 'words_result'-行列表,'words_result_nums'-行数，
    for l1 in ocrJson['words_result']:
        line = OCRLINES.copy()
        line['code'] = l1['words']
        line['coords'] = l1['location']
        line['nums'] = len(l1['words'])
        ocrDict['data'].append(line)
    return ocrDict


# json解析汉王
def ocrJsonHanWangToResult(ocrJson):
    ret = []
    # 全路径文件，文件KB，ocrJson，ocrTxt1，ocrTxt2，ocrPy
    ret.append(ocrJson['name'])
    ret.append(ocrJson['info'][0])
    # 测试-取消原始json
    # ret.append(str(ocrJson))
    ret.append('')
    # 单字符
    lines1 = []         # 单个字组成串
    lineStr0 = []       # 原始字符串
    lineStr1 = []       # 左横向+右横向+纵向字符串简体
    for line in ocrJson['data']:
        # if line['style'] != 0:
        #     # 正文
        #     continue
        if line['type'] == 1:
            # 图像
            continue
        # 标题
        codes = isStrToHz(line['code'])
        if len(codes) <= 0:
            continue
        if len(codes) <= 1:
            lines1.append(line)
        else:
            # 原始文本
            lineStr0.append(codes)
            # 简体文本
            lineStr1.append(strToJianti(codes))
            if line['direction'] == 0:
                # 横向->反向字符串
                lineStr1.append(strToJianti(codes[::-1]))
    # 解析单个字，生成纵向串
    lineStr2 = []  # 上->下字符串
    # 原始
    ocrTxt1 = ''
    for i, s in enumerate(lineStr0):
        if len(ocrTxt1) > 0:
            ocrTxt1 += ' '
        ocrTxt1 += '%d:%s' % (i + 1, s)
    ret.append(ocrTxt1)
    # 简体
    ocrTxt2 = ''
    # 拼音
    ocrPy = ''
    for i, s in enumerate(lineStr1):
        if len(ocrTxt2) > 0:
            ocrTxt2 += ' '
        if len(ocrPy) > 0:
            ocrPy += ' '
        ocrTxt2 += '%d:%s' % (i + 1, s)
        ocrPy += '%d:%s' % (i + 1, strToPinYin(s))
    ret.append(ocrTxt2)
    ret.append(ocrPy)
    return ret


# json解析百度
def ocrJsonBaiduToResult(ocrJson):
    ret = []
    # 全路径文件，文件KB，ocrJson，ocrTxt1，ocrTxt2，ocrPy
    ret.append(ocrJson['name'])
    ret.append(ocrJson['info'][0])
    # 测试-取消原始json
    # ret.append(str(ocrJson['result']))
    ret.append('')
    # 单字符
    lines1 = []         # 单个字组成串
    lineStr0 = []       # 原始字符串
    lineStr1 = []       # 左横向+右横向+纵向字符串简体
    for line in ocrJson['data']:
        # if line['style'] != 0:
        #     # 正文
        #     continue
        # 标题
        codes = isStrToHz(line['code'])
        if len(codes) <= 0:
            continue
        if len(codes) <= 1:
            lines1.append(line)
        else:
            # 原始文本
            lineStr0.append(codes)
            # 简体文本
            lineStr1.append(strToJianti(codes))
            if line['direction'] == 0:
                # 横向->反向字符串
                lineStr1.append(strToJianti(codes[::-1]))
    # 解析单个字，生成纵向串
    # lineStr2 = []  # 上->下字符串
    # 原始
    ocrTxt1 = ''
    for i, s in enumerate(lineStr0):
        if len(ocrTxt1) > 0:
            ocrTxt1 += ' '
        ocrTxt1 += '%d:%s' % (i + 1, s)
    ret.append(ocrTxt1)
    # 简体
    ocrTxt2 = ''
    # 拼音
    ocrPy = ''
    for i, s in enumerate(lineStr1):
        if len(ocrTxt2) > 0:
            ocrTxt2 += ' '
        if len(ocrPy) > 0:
            ocrPy += ' '
        ocrTxt2 += '%d:%s' % (i + 1, s)
        ocrPy += '%d:%s' % (i + 1, strToPinYin(s))
    ret.append(ocrTxt2)
    ret.append(ocrPy)
    return ret


# 识别目录
def ocrDir(dir, save=False):
    url = 'http://%s:%d' % (WEBIP, WEBPORT)
    for d in os.listdir(dir):
        name = pathJoin(dir, d)
        jname = name[:-4] + '.json'
        if not os.path.isfile(name):
            continue
        if name[-4:] not in IMG_EXT:
            continue
        if os.path.isfile(jname):
            # 已经识别
            ocrJson = []
            with open(jname, 'r', encoding="utf-8") as f:
                data = f.read()
                ocrJson = eval(json.loads(data))
        else:
            # 未识别
            result = ocrName(url, name)
            if result['code'] > 0:
                # 出错
                print('识别[%s] 错误：%s' % (name, getOcrState(result['code'])))
                continue
            # 成功
            ocrJson = result['result']
            print('识别[%s] 成功 [宽度-%d,高度-%d,元素数-%d]' %
                  (name, result['Width'], result['Height'], result['Layouts_num']))
            f = open(jname, 'w', encoding="utf-8")
            f.write(json.dumps(str(ocrJson), ensure_ascii=False, indent=4, separators=(',', ':')))
            f.close()
        ocrDict = ocrProc(name, ocrJson)
        ret = ocrJsonToResult(ocrDict)
        print(ret)
        # ret to mysql


if __name__ == '__main__':
    getPinYin()
    getFanTi()
    dir = 'd:/Data/Image'
    # dirName = pathJoin(dir, 'test0')
    # dirName = pathJoin(dir, 'baoding')
    # dirName = pathJoin(dir, 'haidian')
    # dirName = pathJoin(dir, 'paper0')
    dirName = pathJoin(dir, 'test')
    ocrDir(dirName, save=True)
